from src.choice_api.utils import data_str2timestamp,sort_by_first_element


def parse_PHK(data, columns_lst):
    if data is None:
        return None

    company_lst = {}
    # 真实要获取的字段
    total_columns = ["flow", "ph", "cod", "ss", "nh3n", "tp", "tn"]
    if isinstance(columns_lst, str) and columns_lst in total_columns:
        columns_lst = [columns_lst]

    elif isinstance(columns_lst, list) or len(columns_lst) > 0:
        target_correct_columns = []
        for target_column in columns_lst:
            if target_column in total_columns:
                target_correct_columns.append(target_column)
        if len(target_correct_columns) > 0:
            columns_lst = target_correct_columns
        else:
            columns_lst = total_columns
    else:
        columns_lst = total_columns


    param_mapping = {
        "flow": "FLOW",
        "ph": "pH",
        "cod": "COD",
        "ss": "SS",
        "nh3n": "NH3-N",
        "tp": "TP",
        "tn": "TN",
    }

    for row in data:  # 遍历公司
        if 'companyId' not in row:
            continue

        data_lst = []
        for day_data in row['list']:
            data_lst.extend(day_data['list'])

        data_lst2 = []
        for time_data in data_lst:  # 遍历时间
            time_object = {
                "time": time_data['time'],
            }

            if time_data['data']:
                for param in columns_lst:
                    if param in time_data['data']:
                        time_object[param] = time_data['data'][param]
                    else:
                        time_object[param] = -999
            else:
                for param in columns_lst:
                    time_object[param] = -999

            data_lst2.append(time_object)

        company_lst[row['companyName']] = data_lst2

    # 单表csv解析
    # res_columns = ["time"]
    # res_data = [[x["time"]] for x in company_lst[list(company_lst.keys())[0]]]
    #
    # for k, val_lst in company_lst.items():
    #     for param in columns_lst:
    #         res_columns.append(f"{k} {param_mapping[param]}")
    #         for idx, val in enumerate(val_lst):
    #             res_data[idx].append(val[param])
    # return {
    #     "res_columns": res_columns,
    #     "res_data": res_data,
    # }

    # 单表时序指标形式
    res_columns = []
    res_data = {}
    for k, val_lst in company_lst.items():
        for param in columns_lst:
            label = f"{k} 下游排海 {param_mapping[param]}"
            res_columns.append(label)
            res_item_data = []
            for idx, val in enumerate(val_lst):
                if val[param] == -999:
                    continue
                res_item_data.append([data_str2timestamp(val["time"]), val[param]])
            res_data[label] = sort_by_first_element(res_item_data)
    return {
        "res_columns": res_columns,
        "res_data": res_data
    }
